#include <omp.h>

#include "integrate4.h"
#include "common.h"
#include "common_cuda.h"

#define BLOCK_SIZE 512

__global__ void updateVel4(float4* pos, float4* vel, int offset, int n, float dt)
{
	int x = blockIdx.x*BLOCK_SIZE + threadIdx.x + offset;

	if (x < n)
	{
		float4 pos_a = pos[x];
		float3 acc = make_float3(0.0f);

		// sum total force
		int i;
		for (i = 0; i != n; ++i)
		{
			//if (i != x)
			{
				float4 pos_b = pos[i];
				acc += calcGravityAccel(&pos_a, &pos_b);
			}
		}

		// apply acceleration
		vel[x] += make_float4(dt*acc);	// 7 flops
	}
}

__global__ void updatePos4(float4* pos, float4* vel, int offset, int n, float dt)
{
	int x = blockIdx.x*BLOCK_SIZE + threadIdx.x + offset;

	if (x < n)
		pos[x] += dt*vel[x];	// 3 flops
}

void integrate4(float4* d_pos, float4* d_vel, int size, int iterations, float dt)
{
	dim3 dimb(BLOCK_SIZE, 1);
	dim3 dimg((size-1)/dimb.x + 1, 1);

	for (int i = 0; i != iterations; ++i)
	{
		//updateVel4<<<dimg, dimb>>>(d_pos, d_vel, size, dt);
		cudaThreadSynchronize();
		//updatePos<<<dimg, dimb>>>(d_pos, d_vel, size, dt);
		cudaThreadSynchronize();
	}
}

void integrate4(bodyList* bl, int iterations, float dt)
{
	int num_devices = 1;
	cudaGetDeviceCount(&num_devices);

	#pragma omp parallel shared(bl) num_threads(num_devices)
	{
		int tid = omp_get_thread_num();
		cudaSetDevice(tid);

		// allocate memory on gpu
		float4* d_pos = (float4*) safe_cudaMalloc(bl->size*sizeof(float4));
		float4* d_vel = (float4*) safe_cudaMalloc(bl->size*sizeof(float4));

		dim3 dimb(BLOCK_SIZE, 1);
		int total_blocks = (bl->size-1)/BLOCK_SIZE + 1;
		int start_block = (total_blocks*tid)/num_devices;
		int end_block = (total_blocks*(tid+1))/num_devices;
		dim3 dimg(end_block - start_block, 1);
		int offset = start_block*BLOCK_SIZE;
		int size = min(bl->size, end_block*BLOCK_SIZE) - offset;

		// copy to gpu
		for (int i = 0; i != iterations; ++i)
		{
			cudaMemcpy(d_pos, bl->pos, bl->size*sizeof(float4), cudaMemcpyHostToDevice);
			cudaMemcpy(d_vel, bl->vel, bl->size*sizeof(float4), cudaMemcpyHostToDevice);

			updateVel4<<<dimg, dimb>>>(d_pos, d_vel, offset, bl->size, dt);
			cudaThreadSynchronize();
			updatePos4<<<dimg, dimb>>>(d_pos, d_vel, offset, bl->size, dt);
			cudaThreadSynchronize();

			#pragma omp barrier

			// copy results back
			cudaMemcpy(bl->pos + offset, d_pos + offset, size*sizeof(float4), cudaMemcpyDeviceToHost);
			cudaMemcpy(bl->vel + offset, d_vel + offset, size*sizeof(float4), cudaMemcpyDeviceToHost);

			#pragma omp barrier
		}

		// free memory
		cudaFree(d_pos);
		cudaFree(d_vel);
	}
}
